# -*- coding: utf-8 -*-
from pickle import NONE
from .base import BaseTool
from .base import PrintUtils,CmdTask,FileUtils,AptUtils,ChooseTask
from .base import osversion
from .base import run_tool_file
from .base import tr


class RosVersion:
    STATUS_EOL = 0
    STATUS_LTS = 1
    def __init__(self,name,version,status,deps=[]):
        self.name = name
        self.version = version
        self.status = status
        self.deps = deps


class RosVersions:
    ros_version = [
        # ubuntu 18
        RosVersion('melodic', 'ROS1', RosVersion.STATUS_LTS, ['python-catkin-tools','python-rosdep']),
    ]

    @staticmethod
    def get_version_string(name):
        for version in RosVersions.ros_version:
            if version.name == name:
                return "{}({})".format(version.name,version.version)

    @staticmethod
    def get_version(name):
        for version in RosVersions.ros_version:
            if version.name == name:
                return version
       
    @staticmethod
    def install_depend(name):
        depends = RosVersions.get_version(name).deps
        for dep in depends:
            AptUtils.install_pkg(dep)


    @staticmethod
    def tip_test_command(name):
        version = RosVersions.get_version(name).version
        if version=="ROS1":
            PrintUtils.print_warn("提示：可以打开的终端输入roscore测试！")

    @staticmethod
    def get_desktop_version(name):
        version = RosVersions.get_version(name).version
        if version=="ROS1":
            return "ros-{}-desktop-full".format(name)
        elif version=="ROS2":
            return "ros-{}-desktop".format(name)


key_urls = [
    'https://gitee.com/ohhuo/rosdistro/raw/master/ros.asc',
    'https://raw.githubusercontent.com/ros/rosdistro/master/ros.asc',
]

ros_mirror_dic = {
    "tsinghua":{"ROS1":"http://mirrors.tuna.tsinghua.edu.cn/ros/ubuntu/"},
    "huawei":{"ROS1":"https://repo.huaweicloud.com/ros/ubuntu/"},
    "packages.ros":{"ROS1":"http://packages.ros.org/ros/ubuntu/"},
    "https.packages.ros":{"ROS1":"https://packages.ros.org/ros/ubuntu/"},
}


ros_dist_dic = {
    'artful':{"packages.ros"},
    'bionic':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'buster':{"packages.ros"},
    'cosmic':{"packages.ros"},
    'disco':{"packages.ros"},
    'eoan':{"packages.ros"},
    'focal':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'jessie':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'lucid':{"packages.ros"},
    'maverick':{"packages.ros"},
    'natty':{"packages.ros"},
    'oneiric':{"packages.ros"},
    'precise':{"packages.ros"},
    'quantal':{"packages.ros"},
    'raring':{"packages.ros"},
    'saucy':{"packages.ros"},
    'stretch':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'trusty':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'utopic':{"packages.ros"},
    'vivid':{"packages.ros"},
    'wheezy':{"packages.ros"},
    'wily':{"packages.ros"},
    'xenial':{"tsinghua","huawei","packages.ros","https.packages.ros"},
    'yakkety':{"packages.ros"},
    'zesty':{"packages.ros"},
}




class Tool(BaseTool):
    def __init__(self):
        self.name = "一键安装ROS Jetson"
        self.type = BaseTool.TYPE_INSTALL
        self.author = '小鱼'


    def get_mirror_by_code(self,code,arch='amd64',first_choose="tsinghua"):
        """
        获取镜像通过系统版本号
        """
        ros1_choose_queue = [first_choose,"tsinghua","huawei","packages.ros"]
        
 

        mirror = []
        # 确认源里有对应的系统的，比如jammy
        if code in ros_dist_dic.keys():
            for item in ros1_choose_queue:
                if item in ros_dist_dic[code]:
                    mirror.append(ros_mirror_dic[item]['ROS1'])
                    break
        # if code in ['focal']:
        #     mirror.append(ros_mirror_dic['packages.ros']['ROS2'])
        return mirror


    def add_key(self):
        PrintUtils.print_success(tr.tr('============正在添加ROS源密钥================='))
        # check apt
        if not AptUtils.checkapt(): 
            pass
        # install dep
        AptUtils.install_pkg('curl')
        AptUtils.install_pkg('gnupg2')

        # add key
        PrintUtils.print_success(tr.tr('正在挑选最快的密钥服务:{}').format(key_urls))
        key_url = AptUtils.get_fast_url(key_urls)
        if not key_url: 
            PrintUtils.print_error(tr.tr("获取密钥失败"))
            return
        key_url = key_url[0]
        PrintUtils.print_success(tr.tr('已自动选择最快密钥服务:{}').format(key_url))

        cmd_result = CmdTask("curl -s {} | sudo apt-key add -".format(key_url)).run()
        if cmd_result[0]!=0: 
            cmd_result = CmdTask("curl -s {} | sudo apt-key add -".format(key_url)).run()
            
        if cmd_result[0]!=0:
            PrintUtils.print_info(tr.tr("导入密钥失败，开始更换导入方式并二次尝试..."))
            cmd_result = CmdTask("sudo apt-key adv --keyserver keyserver.ubuntu.com --recv-keys F42ED6FBAB17C654",10).run()

        # 针对trusted.gpg.d问题解决方案
        if FileUtils.check_result(cmd_result,['trusted.gpg.d']):
            cmd_result = CmdTask("curl -s {} | sudo gpg --no-default-keyring --keyring gnupg-ring:/etc/apt/trusted.gpg.d/ros.gpg --import".format(key_url)).run()
            cmd_result = CmdTask("sudo chmod 644 /etc/apt/trusted.gpg.d/ros.gpg",10).run()

        return cmd_result


    def check_sys_source(self):
        # 更换系统源
        PrintUtils.print_warn("=========接下来这一步很很很很重要，如果不知道怎么选请选择1========")
        
        # 直接选择更换系统源
        tool = run_tool_file('tools.tool_config_system_source', authorun=False)
        tool.change_sys_source()


    def get_all_instsll_ros_pkgs(self):
        AptUtils.checkapt()
        dic_base = AptUtils.search_package('ros-base','ros-[A-Za-z]+-ros-base',"ros-","-base")
        if dic_base== None: return None
        ros_name = {}
        for a in dic_base.keys(): 
            ros_name[RosVersions.get_version_string(a)] = a
        if len(ros_name) == 0:
            return None
        return ros_name

    def add_source(self):
        """
        检查并添加ROS系统源
        """
        arch = AptUtils.getArch()
        if arch==None: return False

        #add source 1
        mirrors = self.get_mirror_by_code(osversion.get_codename(),arch=arch)
        PrintUtils.print_info("根据您的系统，为您推荐安装源为{}".format(mirrors))
        source_data = ''
        for mirror in mirrors:
            source_data += 'deb [arch={}]  {} {} main\n'.format(arch,mirror,osversion.get_codename())
        FileUtils.delete('/etc/apt/sources.list.d/ros-fish.list')
        FileUtils.new('/etc/apt/sources.list.d/',"ros-fish.list",source_data)

        ros_pkg = self.get_all_instsll_ros_pkgs()
        if ros_pkg and len(ros_pkg)>1:
            PrintUtils.print_success("成功添加ROS源，安装ROS！") 
            return
        
        #add source2 
        PrintUtils.print_warn("换源后更新失败，第二次开始切换源，尝试更换ROS2源为华为源！") 
        mirrors = self.get_mirror_by_code(osversion.get_codename(),arch=arch,first_choose="huawei")
        PrintUtils.print_info("根据您的系统，为您推荐安装源为{}".format(mirrors))
        source_data = ''
        for mirror in mirrors:
            source_data += 'deb [arch={}]  {} {} main\n'.format(arch,mirror,osversion.get_codename())
        FileUtils.delete('/etc/apt/sources.list.d/ros-fish.list')
        FileUtils.new('/etc/apt/sources.list.d/',"ros-fish.list",source_data)
        ros_pkg = self.get_all_instsll_ros_pkgs()
        if ros_pkg and len(ros_pkg)>1:
            PrintUtils.print_success("成功添加ROS源，安装ROS") 
            return


        PrintUtils.print_warn("换源后更新失败，第三次开始切换源，尝试使用https-ROS官方源") 
        mirrors = self.get_mirror_by_code(osversion.get_codename(),arch=arch,first_choose="https.packages.ros")
        PrintUtils.print_info("根据您的系统，为您推荐安装源为{}".format(mirrors))
        source_data = ''
        for mirror in mirrors:
            source_data += 'deb [arch={}]  {} {} main\n'.format(arch,mirror,osversion.get_codename())
        FileUtils.delete('/etc/apt/sources.list.d/ros-fish.list')
        FileUtils.new('/etc/apt/sources.list.d/',"ros-fish.list",source_data)
        ros_pkg = self.get_all_instsll_ros_pkgs()
        if ros_pkg and len(ros_pkg)>1:
            PrintUtils.print_success("成功添加ROS源，安装ROS") 
            return

        #add source2 
        PrintUtils.print_warn("换源后更新失败，第四次开始切换源，尝试更换ROS源为http-ROS官方源") 
        mirrors = self.get_mirror_by_code(osversion.get_codename(),arch=arch,first_choose="packages.ros")
        PrintUtils.print_info("根据您的系统，为您推荐安装源为{}".format(mirrors))
        source_data = ''
        for mirror in mirrors:
            source_data += 'deb [arch={}]  {} {} main\n'.format(arch,mirror,osversion.get_codename())
        FileUtils.delete('/etc/apt/sources.list.d/ros-fish.list')
        FileUtils.new('/etc/apt/sources.list.d/',"ros-fish.list",source_data)
        ros_pkg = self.get_all_instsll_ros_pkgs()
        if ros_pkg and len(ros_pkg)>1:
            PrintUtils.print_success("成功添加ROS源，安装ROS") 
            return

        # echo >>/etc/apt/apt.conf.d/99verify-peer.conf "Acquire { https::Verify-Peer false }"
        if  not AptUtils.checkapt(): PrintUtils.print_error("四次换源后都失败！") 


    def support_install(self):
        # check support
        if (osversion.get_codename() not in ros_dist_dic.keys()) and (osversion.get_codename() not in ros2_dist_dic.keys()):
            PrintUtils.print_error("小鱼:检测当前系统{}{}:{}".format(osversion.get_name(), osversion.get_version(),osversion.get_codename()))
            return False
        PrintUtils.print_success("小鱼:检测当前系统{}{}:{}".format(osversion.get_name(), osversion.get_version(),osversion.get_codename()))
        return True

    def install_success(self,name):
        """
        检查某个版本的ROS是否安装成功
        """
        result = CmdTask("ls /opt/ros/{}/setup.bash".format(name), 0).run()
        if str(result[1]).find('setup.bash') >= 1:
            return True
        return False

    

def choose_and_install_ros(self):
    # search ros packages
    # 直接选择melodic版本
    rosname = 'melodic'
    install_version = rosname

    version_dic = {1: rosname + "桌面版", 2: rosname + "基础版(小)"}
    code = 2  # 直接选择基础版(小)

    install_tool = 'aptitude'
    install_tool_apt = 'apt'
    if osversion.get_version() == "16.04":
        install_tool = 'apt'

    if install_tool == 'aptitude':
        AptUtils.install_pkg('aptitude')

    # ======================================================基础版本==============================================================
    if code == 2:
        # 第一次尝试
        cmd_result = CmdTask("sudo {} install {} -y".format(install_tool_apt, dic_base[install_version]), 300, os_command=True).run()
        cmd_result = CmdTask("sudo {} install {} -y".format(install_tool_apt, dic_base[install_version]), 300, os_command=False).run()
        if FileUtils.check_result(cmd_result, ['未满足的依赖关系', 'unmet dependencies', 'but it is not installable']):
            # 尝试使用aptitude解决依赖问题
            PrintUtils.print_warn("============================================================")
            PrintUtils.print_delay("出现依赖问题，请稍后输入n,再选择y")
            # 移除确认输入，直接安装
            cmd_result = CmdTask("sudo {} install   {}".format(install_tool, dic_base[install_version]), 300, os_command=True).run()
            cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool, dic_base[install_version]), 300, os_command=False).run()

    # ======================================================桌面版本==============================================================
    elif code == 1:
        cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool_apt, RosVersions.get_desktop_version(install_version)), 300, os_command=True).run()
        cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool_apt, RosVersions.get_desktop_version(install_version)), 300, os_command=False).run()
        if FileUtils.check_result(cmd_result, ['未满足的依赖关系', 'unmet dependencies', 'but it is not installable']):
            # 尝试使用aptitude解决依赖问题
            PrintUtils.print_warn("============================================================")
            PrintUtils.print_delay("出现依赖问题，请稍后输入n,再选择y,稍后手动运行命令: sudo aptitude install {})".format(RosVersions.get_desktop_version(install_version)))
            # 移除确认输入，直接安装
            cmd_result = CmdTask("sudo {} install   {}".format(install_tool, RosVersions.get_desktop_version(install_version)), 300, os_command=True).run()
            cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool, RosVersions.get_desktop_version(install_version)), 300, os_command=False).run()

    # apt broken error
    if cmd_result[0] != 0:
        if FileUtils.check_result(cmd_result, ['--fix-broken install']):
            CmdTask("sudo apt --fix-broken install -y").run()
            if code == 2:
                cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool, dic_base[rosname]), 300, os_command=False).run()
            elif code == 1:
                cmd_result = CmdTask("sudo {} install   {} -y".format(install_tool, rosname), 300, os_command=False).run()

    # 安装额外的依赖
    RosVersions.install_depend(install_version)

    return install_version


    def config_env_and_tip(self,version_name):
        """
        配置系统源
        """
        if self.install_success(version_name):
            run_tool_file('tools.tool_config_rosenv')
            PrintUtils.print_success("安装成功")
        else: PrintUtils.print_error("安装失败")


    def install_ros(self):
        if not self.support_install(): return False

        self.check_sys_source()
        self.add_key()
        self.add_source()

        ros_version = self.choose_and_install_ros()
        if not ros_version:
            return
        self.config_env_and_tip(ros_version)
        if self.install_success(ros_version):
            RosVersions.tip_test_command(ros_version)


    def run(self):
        self.install_ros()

